# 环节5：混合检索（Elasticsearch关键词 + 语义向量）
# 功能：结合ES的关键词精确匹配和向量的语义相似性，提高检索准确性
# 依赖：需安装 elasticsearch==8.14.0，确保ES服务和Milvus服务已启动

import time
from typing import List, Dict
from elasticsearch import Elasticsearch, exceptions
from pymilvus import Collection, connections, utility
from langchain_community.embeddings import DashScopeEmbeddings

# 配置参数
CONFIG = {
    # 通义千问配置
    "tongyi_api_key": "sk-6434b554122c4b049ceb805d703f695b",
    # Elasticsearch配置
    "es_host": "localhost",
    "es_port": 9200,
    "es_index": "insurance_manual",
    # Milvus配置
    "milvus_host": "localhost",
    "milvus_port": "19530",
    "milvus_collection": "insurance_manual",
    # 混合检索参数
    "top_k_es": 5,  # ES关键词检索返回数量
    "top_k_semantic": 5,  # 语义检索返回数量
    "es_weight": 0.4,  # 关键词检索权重
    "semantic_weight": 0.6  # 语义检索权重
}


class HybridRetriever:
    def __init__(self):
        # 初始化组件
        self.embeddings = self._init_embeddings()
        self.es = self._init_es()
        self.milvus_collection = self._init_milvus()
        self.vector_dim = self._get_vector_dim()

    def _init_embeddings(self) -> DashScopeEmbeddings:
        """初始化通义千问嵌入模型"""
        try:
            embeddings = DashScopeEmbeddings(
                model="text-embedding-v4",
                dashscope_api_key=CONFIG["tongyi_api_key"]
            )
            # 验证模型
            embeddings.embed_query("测试")
            print("✅ 通义千问嵌入模型初始化成功")
            return embeddings
        except Exception as e:
            raise Exception(f"嵌入模型初始化失败: {str(e)}")

    def _init_es(self) -> Elasticsearch:
        """初始化Elasticsearch连接"""
        try:
            es = Elasticsearch(
                hosts=[f"http://{CONFIG['es_host']}:{CONFIG['es_port']}"],
                timeout=30
            )
            # 验证连接
            if not es.ping():
                raise exceptions.ConnectionError("无法连接到Elasticsearch")

            # 检查索引是否存在，不存在则创建
            if not es.indices.exists(index=CONFIG["es_index"]):
                es.indices.create(
                    index=CONFIG["es_index"],
                    mappings={
                        "properties": {
                            "text": {"type": "text", "analyzer": "ik_max_word"},  # 中文分词
                            "page": {"type": "integer"},
                            "id": {"type": "keyword"}
                        }
                    }
                )
                print(f"✅ 创建ES索引: {CONFIG['es_index']}")
            else:
                print(f"✅ 连接到ES索引: {CONFIG['es_index']}")
            return es
        except Exception as e:
            raise Exception(f"ES初始化失败: {str(e)}")

    def _init_milvus(self) -> Collection:
        """初始化Milvus连接"""
        try:
            if not connections.has_connection("default"):
                connections.connect(
                    alias="default",
                    host=CONFIG["milvus_host"],
                    port=CONFIG["milvus_port"]
                )

            if not utility.has_collection(CONFIG["milvus_collection"]):
                raise Exception(f"Milvus集合不存在: {CONFIG['milvus_collection']}")

            collection = Collection(CONFIG["milvus_collection"])
            collection.load()
            print(f"✅ 加载Milvus集合: {CONFIG['milvus_collection']}")
            return collection
        except Exception as e:
            raise Exception(f"Milvus初始化失败: {str(e)}")

    def _get_vector_dim(self) -> int:
        """获取向量维度"""
        vector_field = next(f for f in self.milvus_collection.schema.fields if f.name == "vector")
        return vector_field.params["dim"]

    def es_keyword_search(self, query: str) -> List[Dict]:
        """Elasticsearch关键词检索"""
        try:
            # 构建查询：匹配text字段，使用中文分词
            response = self.es.search(
                index=CONFIG["es_index"],
                size=CONFIG["top_k_es"],
                query={
                    "match": {
                        "text": {
                            "query": query,
                            "analyzer": "ik_smart"  # 智能分词，提高精度
                        }
                    }
                }
            )

            # 处理结果
            results = []
            for hit in response["hits"]["hits"]:
                results.append({
                    "id": hit["_source"]["id"],
                    "text": hit["_source"]["text"],
                    "page": hit["_source"]["page"],
                    "score": hit["_score"],  # ES的相关性评分
                    "type": "keyword"
                })
            return results
        except Exception as e:
            print(f"ES检索警告: {str(e)}")
            return []

    def semantic_search(self, query: str) -> List[Dict]:
        """Milvus语义向量检索"""
        try:
            # 生成查询向量
            query_vector = self.embeddings.embed_query(query)

            # 验证向量维度
            if len(query_vector) != self.vector_dim:
                raise ValueError(f"向量维度不匹配: 预期{self.vector_dim}, 实际{len(query_vector)}")

            # 语义检索
            search_params = {
                "metric_type": "L2",
                "params": {"nprobe": 10}
            }

            results = self.milvus_collection.search(
                data=[query_vector],
                anns_field="vector",
                param=search_params,
                limit=CONFIG["top_k_semantic"],
                output_fields=["text", "page"]
            )

            # 处理结果（转换L2距离为相似度）
            processed = []
            for hit in results[0]:
                # L2距离越小越相似，转换为0-1之间的相似度
                similarity = 1 / (1 + hit.distance)
                processed.append({
                    "id": hit.id,
                    "text": hit.entity.text,
                    "page": hit.entity.page,
                    "score": similarity,  # 语义相似度
                    "type": "semantic"
                })
            return processed
        except Exception as e:
            print(f"语义检索警告: {str(e)}")
            return []

    def hybrid_search(self, query: str) -> List[Dict]:
        """混合检索：融合关键词和语义检索结果"""
        start_time = time.time()

        # 1. 并行执行两种检索
        es_results = self.es_keyword_search(query)
        semantic_results = self.semantic_search(query)

        # 2. 标准化评分（将不同来源的评分映射到0-1范围）
        if es_results:
            max_es_score = max(item["score"] for item in es_results)
            for item in es_results:
                item["normalized_score"] = item["score"] / max_es_score * CONFIG["es_weight"]

        if semantic_results:
            max_sem_score = max(item["score"] for item in semantic_results)
            for item in semantic_results:
                item["normalized_score"] = item["score"] / max_sem_score * CONFIG["semantic_weight"]

        # 3. 合并结果（去重，保留最高分）
        merged = {}
        for item in es_results + semantic_results:
            if item["id"] not in merged:
                merged[item["id"]] = item
            else:
                # 同一文档取更高的标准化评分
                if item["normalized_score"] > merged[item["id"]]["normalized_score"]:
                    merged[item["id"]] = item

        # 4. 按标准化评分排序
        final_results = sorted(
            merged.values(),
            key=lambda x: x["normalized_score"],
            reverse=True
        )

        # 5. 输出检索统计
        print(f"\n🔍 混合检索完成（耗时{time.time() - start_time:.2f}秒）")
        print(f"   问题: {query}")
        print(f"   关键词检索结果: {len(es_results)}条")
        print(f"   语义检索结果: {len(semantic_results)}条")
        print(f"   去重后合并结果: {len(final_results)}条")

        return final_results

    def add_to_es(self, docs: List[Dict]):
        """将文档添加到ES（用于初始化数据，实际应在环节4与Milvus同步插入）"""
        try:
            for doc in docs:
                self.es.index(
                    index=CONFIG["es_index"],
                    id=doc["id"],
                    document={
                        "text": doc["text"],
                        "page": doc["page"],
                        "id": doc["id"]
                    }
                )
            self.es.indices.refresh(index=CONFIG["es_index"])
            print(f"✅ 已向ES索引添加{len(docs)}条文档")
        except Exception as e:
            raise Exception(f"ES数据插入失败: {str(e)}")


def display_results(results: List[Dict]):
    """格式化显示检索结果"""
    for i, doc in enumerate(results[:3], 1):  # 只显示前3条最相关结果
        print(f"\n【混合检索结果 {i}】")
        print(f"页码: 第{doc['page']}页")
        print(f"综合评分: {doc['normalized_score']:.4f}")
        print(f"来源: {doc['type']}检索")
        print(f"内容: {doc['text'][:200]}...")  # 显示前200字符
        print("-" * 100)


def main():
    """混合检索主函数"""
    try:
        # 初始化混合检索器
        retriever = HybridRetriever()

        # 测试问题列表
        test_queries = [
            "平安福2024版的投保条件是什么？",
            "健康告知需要提供哪些信息？",
            "保险理赔的流程是怎样的？",
            "平安福的缴费方式有哪些？"
        ]

        # 执行混合检索并显示结果
        for query in test_queries:
            results = retriever.hybrid_search(query)
            display_results(results)

        print("\n🎉 混合检索流程完成")

    except Exception as e:
        print(f"\n❌ 检索失败: {str(e)}")


if __name__ == "__main__":
    main()
